import logging
import traceback
import warnings
from urllib.parse import unquote

from flask import Flask, request
from gevent import pywsgi
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

warnings.filterwarnings('ignore')

# 放弃向huggingface联网访问。
# os.environ['TRANSFORMERS_OFFLINE'] = "1"
# os.environ['HF_HUB_OFFLINE'] = "1"

# 初始化日志
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# 创建文件处理器，将日志输出到文件
file_handler = logging.FileHandler('app.log')
file_handler.setLevel(logging.DEBUG)
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(file_formatter)

# 创建控制台处理器，将日志输出到控制台
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
console_handler.setFormatter(console_formatter)

# 添加处理器到日志器
logger.addHandler(file_handler)
logger.addHandler(console_handler)

app = Flask(__name__)
app.logger = logger

# 模型配置
models = {
    # 中文 ==> 外文
    'zh_en': ('Helsinki-NLP/opus-mt-zh-en', 'translation_zh_to_en'),
    # 中译日 效果不好
    # 'zh_ja': ('Helsinki-NLP/opus-mt-tc-big-zh-ja', 'translation_zh_to_ja'),
    'zh_de': ('Helsinki-NLP/opus-mt-zh-de', 'translation_zh_to_de'),
    # 英文 ==> 中文、外文
    'en_zh': ('Helsinki-NLP/opus-mt-en-zh', 'translation_en_to_zh'),
    # 日文 ==> 中文、外文
    'ja_zh': ('neverLife/nllb-200-distilled-600M-ja-zh', 'translation_ja_to_zh'),
    'ja_en': ('Helsinki-NLP/opus-mt-ja-en', 'translation_ja_to_en'),
    # 德文 ==> 中文、外文
    'de_zh': ('Helsinki-NLP/opus-mt-de-ZH', 'translation_de_to_zh'),
}

# 模型缓存，用于懒加载
translations = {}


def load_translation_model(tran_type):
    """懒加载翻译模型"""
    if tran_type not in translations:
        model_name, task = models[tran_type]
        app.logger.info(f"load model [{model_name}], for [{task}].")
        if tran_type == 'ja_zh':
            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang="jpn_Jpan", tgt_lang="zho_Hans")
            translations[tran_type] = (model, tokenizer)
        else:
            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            translations[tran_type] = pipeline(task, model=model, tokenizer=tokenizer)
        app.logger.info(f"load model [{model_name}], for [{task}] success.")


def translate(text, tran_type):
    app.logger.info(f"Received request: type = {tran_type}, text = {text}")
    result = ''
    if not text:
        return result

    try:
        load_translation_model(tran_type)
        if tran_type == 'ja_zh':
            model, tokenizer = translations[tran_type]
            input_ids = tokenizer.encode(text, max_length=512, padding=True, return_tensors='pt')
            outputs = model.generate(input_ids, num_beams=4, max_new_tokens=128)
            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        else:
            translation_pipeline = translations[tran_type]
            result = translation_pipeline(text, max_length=512)[0]['translation_text']

        app.logger.info(f"Translation result = {result}")
    except Exception as e:
        app.logger.error(f"Translation failed: {e}")
        app.logger.debug(traceback.format_exc())

    return result


@app.route('/tr', methods=['GET'])
def tr_get():
    text = request.args.get('from', '')
    tran_type = request.args.get('type', 'zh_en')
    return translate(unquote(text), tran_type)


@app.route('/tr', methods=['POST'])
def tr_post():
    text = request.form.get('from', '')
    tran_type = request.form.get('type', 'zh_en')
    return translate(text, tran_type)


if __name__ == '__main__':
    app.logger.info(f"server starting...")
    # 禁用 flask内置 werkzeug 日志输出
    server = pywsgi.WSGIServer(('0.0.0.0', 20000), app, log=None)
    app.logger.info(f"server start up.")
    server.serve_forever()
